import pandas as pd
from rouge import Rouge
from nltk.translate.bleu_score import SmoothingFunction
from nltk.translate.bleu_score import sentence_bleu

top_distinct_num = 10


def deduplicate_by_threshold(origin_texts, candidates, similarity_threshold):
    rouge = Rouge()
    # cc = SmoothingFunction()
    for idx, text in enumerate(origin_texts):
        if text == "":
            continue
        hasDeduplication = False
        for can in candidates:
            similarity = rouge.get_scores(text, can[1])[0]['rouge-l']['f']
            # similarity = sentence_bleu([can[1].split()], text.split(), smoothing_function=cc.method4)
            if similarity >= similarity_threshold:
                hasDeduplication = True
                break

        if not hasDeduplication:
            candidates.append([idx, text])
            if len(candidates) >= top_distinct_num:
                break


# pass n texts to get top-10 distinct texts
def deduplicator(origin_texts, initial_similarity_threshold, keep_searched=True):
    similarity_threshold = initial_similarity_threshold
    candidates = []
    while len(candidates) != top_distinct_num and similarity_threshold <= 1.0:
        if not keep_searched:
            candidates = []
        deduplicate_by_threshold(origin_texts, candidates, similarity_threshold)
        similarity_threshold += 0.1

    temp = []
    for can in candidates:
        temp.append("[{}]{}".format(can[0], can[1]))
    candidates = temp
    return candidates


predictions = pd.read_csv('original-model-hypothesis/SOTitle-bos-eos-1063-30/predictions-bs-100-lower.csv')
predictions.fillna("", inplace=True)
predictions = predictions.values.tolist()
init_simi_threshold = 0.0
while init_simi_threshold <= 1.0:
    deduplicated_predictions = []
    for idx, pres in enumerate(predictions):
        deduplicated_predictions.append(deduplicator(pres[1:101], init_simi_threshold, False))
        if idx % 100 == 0:
            print('Progress[{}/{}]'.format(idx, len(predictions)), deduplicated_predictions[-1][:5])
    deduplicated_predictions = pd.DataFrame(deduplicated_predictions)
    deduplicated_predictions.to_csv(
        'original-model-hypothesis/SOTitle-bos-eos-1063-30/deduplication-rouge/hypothesis-bs-top10-clearsearched-rouge'
        + str(init_simi_threshold) + '.csv')

    deduplicated_predictions = []
    for idx, pres in enumerate(predictions):
        deduplicated_predictions.append(deduplicator(pres[1:101], init_simi_threshold, True))
        if idx % 100 == 0:
            print('Progress[{}/{}]'.format(idx, len(predictions)), deduplicated_predictions[-1][:5])
    deduplicated_predictions = pd.DataFrame(deduplicated_predictions)
    deduplicated_predictions.to_csv(
        'original-model-hypothesis/SOTitle-bos-eos-1063-30/deduplication-rouge/hypothesis-bs-top10-keepsearched-rouge'
        + str(init_simi_threshold) + '.csv')
    init_simi_threshold += 0.1
